Greptile Summary

This PR adds selective CPU activation offloading to the SReLU fused op (_ForwardGroupedMLP_CuTeGEMMBase_MXFP8). In the pre-quantized MXFP8 input path the FC1 input is now wrapped in a GroupedTensorStorage (previously used directly as GroupedTensor), and the GroupedTensorStorage type is also used for the FC2 GEMM input. A new get_data_tensors() method on GroupedTensorStorage enables the V1 offload path to enumerate component tensors without destructive side-effects.

• FC1 input and activation-function input are selectively offloaded based on fine_grained_activation_offloading on the respective ops; weights are always pinned to GPU.
• FC2 GEMM input (saved_grouped_fc2_x) is always pinned to GPU with mark_not_offload, without a corresponding fine_grained_activation_offloading check, making selective offloading of that tensor impossible through the new API.
• In the non-V1 offload path, mark_not_offload internally calls prepare_for_saving/restore_from_saved on GroupedTensorStorage objects, which clears quantized_tensors as a side effect and does not restore it; this works safely today because those objects are freshly constructed, but the pattern is fragile for future callers.

Confidence Score: 4/5

The selective offloading logic is functional for FC1 input and activation tensors; the main gap is that the FC2 GEMM input has no corresponding selective-offload control and is always pinned to GPU.

The core mechanism works correctly end-to-end. The FC2 input is unconditionally prevented from offloading with no API to override this, creating an asymmetry that may be an oversight. The non-V1 path uses prepare_for_saving destructively on GroupedTensorStorage objects but is safe here because quantized_tensors is always None at the point of the call.

forward_grouped_mlp.py around the FC2 mark_not_offload block at lines 596-599 where the missing fine_grained_activation_offloading gate lives

Important Files Changed

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Forward fused op] --> B[Construct GroupedTensorStorage for fc1_x and fc2_x]
    B --> C{cpu_offloading enabled?}
    C -- No --> G[save_for_backward via OperationContext]
    C -- Yes --> D{offload_fc1_input?}
    D -- False --> E[mark_not_offload grouped_fc1_x]
    D -- True --> F[fc1_x eligible for offload]
    E --> H[mark_not_offload fc1 weights always]
    F --> H
    H --> I{offload_activation_input?}
    I -- False --> J[mark_not_offload activation_in and scales]
    I -- True --> K[activation tensors eligible for offload]
    J --> L[mark_not_offload saved_grouped_fc2_x always]
    K --> L
    L --> M[mark_not_offload fc2 weights always]
    M --> G
    G --> N[fuser.py prepare_for_saving decomposes objects]
    N --> O[PyTorch save_for_backward push_tensor hook]
    O --> P{_TE_do_not_offload set?}
    P -- Yes --> Q[Stay on GPU]
    P -- No --> R[Offload to CPU]

_{Reviews (4): Last reviewed commit: "Simplify fused grouped MLP offload check..." | Re-trigger Greptile}